"""
特殊矩阵处理模块
包含稀疏矩阵、带状矩阵、Toeplitz矩阵、循环矩阵、Hankel矩阵等特殊结构矩阵的处理
支持高效存储、快速运算和结构保持的算法
"""
import numpy as np
from scipy import sparse
from scipy.linalg import toeplitz, hankel, circulant
from typing import Union, Optional, Tuple, List
import warnings

class SpecialMatrices:
    """特殊矩阵处理类"""
    
    # ==================== 稀疏矩阵处理 ====================
    
    @staticmethod
    def create_sparse_coo(data: np.ndarray, row: np.ndarray, col: np.ndarray, 
                         shape: Tuple[int, int]) -> sparse.coo_matrix:
        """
        创建COO格式稀疏矩阵
        :param data: 非零元素值
        :param row: 行索引
        :param col: 列索引
        :param shape: 矩阵形状
        :return: COO稀疏矩阵
        """
        return sparse.coo_matrix((data, (row, col)), shape=shape)
    
    @staticmethod
    def create_sparse_csr(data: np.ndarray, indices: np.ndarray, 
                         indptr: np.ndarray, shape: Tuple[int, int]) -> sparse.csr_matrix:
        """
        创建CSR格式稀疏矩阵
        :param data: 非零元素值
        :param indices: 列索引
        :param indptr: 行指针
        :param shape: 矩阵形状
        :return: CSR稀疏矩阵
        """
        return sparse.csr_matrix((data, indices, indptr), shape=shape)
    
    @staticmethod
    def create_sparse_csc(data: np.ndarray, indices: np.ndarray, 
                         indptr: np.ndarray, shape: Tuple[int, int]) -> sparse.csc_matrix:
        """
        创建CSC格式稀疏矩阵
        :param data: 非零元素值
        :param indices: 行索引
        :param indptr: 列指针
        :param shape: 矩阵形状
        :return: CSC稀疏矩阵
        """
        return sparse.csc_matrix((data, indices, indptr), shape=shape)
    
    @staticmethod
    def convert_sparse_format(matrix: sparse.spmatrix, target_format: str) -> sparse.spmatrix:
        """
        转换稀疏矩阵格式
        :param matrix: 输入稀疏矩阵
        :param target_format: 目标格式 ('coo', 'csr', 'csc', 'lil', 'dok')
        :return: 转换后的稀疏矩阵
        """
        format_map = {
            'coo': matrix.tocoo,
            'csr': matrix.tocsr,
            'csc': matrix.tocsc,
            'lil': matrix.tolil,
            'dok': matrix.todok
        }
        
        if target_format not in format_map:
            raise ValueError(f"Unsupported format: {target_format}")
        
        return format_map[target_format]()
    
    @staticmethod
    def sparse_matrix_stats(matrix: sparse.spmatrix) -> dict:
        """
        稀疏矩阵统计信息
        :param matrix: 稀疏矩阵
        :return: 统计信息字典
        """
        return {
            'shape': matrix.shape,
            'nnz': matrix.nnz,
            'density': matrix.nnz / (matrix.shape[0] * matrix.shape[1]),
            'format': matrix.format,
            'dtype': matrix.dtype,
            'has_sorted_indices': matrix.has_sorted_indices if hasattr(matrix, 'has_sorted_indices') else None
        }
    
    # ==================== 带状矩阵处理 ====================
    
    @staticmethod
    def create_band_matrix(diagonals: List[np.ndarray], offsets: List[int], 
                          shape: Tuple[int, int]) -> sparse.dia_matrix:
        """
        创建带状矩阵
        :param diagonals: 对角线数据列表
        :param offsets: 对角线偏移量列表
        :param shape: 矩阵形状
        :return: 带状矩阵
        """
        return sparse.diags(diagonals, offsets, shape=shape, format='dia')
    
    @staticmethod
    def extract_band_matrix(matrix: np.ndarray, lower_bandwidth: int, 
                           upper_bandwidth: int) -> np.ndarray:
        """
        提取矩阵的带状部分
        :param matrix: 输入矩阵
        :param lower_bandwidth: 下带宽
        :param upper_bandwidth: 上带宽
        :return: 带状矩阵
        """
        m, n = matrix.shape
        result = np.zeros_like(matrix)
        
        for i in range(m):
            for j in range(n):
                if -lower_bandwidth <= j - i <= upper_bandwidth:
                    result[i, j] = matrix[i, j]
        
        return result
    
    @staticmethod
    def band_matrix_solve(band_matrix: sparse.dia_matrix, b: np.ndarray) -> np.ndarray:
        """
        求解带状矩阵线性方程组
        :param band_matrix: 带状矩阵
        :param b: 右端向量
        :return: 解向量
        """
        return sparse.linalg.spsolve(band_matrix.tocsr(), b)
    
    # ==================== Toeplitz矩阵 ====================
    
    @staticmethod
    def create_toeplitz_matrix(c: np.ndarray, r: Optional[np.ndarray] = None) -> np.ndarray:
        """
        创建Toeplitz矩阵
        :param c: 第一列
        :param r: 第一行（可选）
        :return: Toeplitz矩阵
        """
        return toeplitz(c, r)
    
    @staticmethod
    def toeplitz_matrix_vector_multiply(c: np.ndarray, r: np.ndarray, 
                                       x: np.ndarray) -> np.ndarray:
        """
        Toeplitz矩阵与向量的快速乘法（使用FFT）
        :param c: Toeplitz矩阵第一列
        :param r: Toeplitz矩阵第一行
        :param x: 向量
        :return: 乘积结果
        """
        from scipy.linalg import solve_toeplitz
        # 这里简化实现，实际应使用FFT优化
        T = toeplitz(c, r)
        return T @ x
    
    @staticmethod
    def solve_toeplitz_system(c: np.ndarray, b: np.ndarray) -> np.ndarray:
        """
        求解Toeplitz线性方程组
        :param c: Toeplitz矩阵第一列
        :param b: 右端向量
        :return: 解向量
        """
        from scipy.linalg import solve_toeplitz
        return solve_toeplitz(c, b)
    
    # ==================== 循环矩阵 ====================
    
    @staticmethod
    def create_circulant_matrix(c: np.ndarray) -> np.ndarray:
        """
        创建循环矩阵
        :param c: 第一行/列
        :return: 循环矩阵
        """
        return circulant(c)
    
    @staticmethod
    def circulant_eigenvalues(c: np.ndarray) -> np.ndarray:
        """
        计算循环矩阵的特征值（使用DFT）
        :param c: 循环矩阵第一行
        :return: 特征值
        """
        return np.fft.fft(c)
    
    @staticmethod
    def circulant_matrix_vector_multiply(c: np.ndarray, x: np.ndarray) -> np.ndarray:
        """
        循环矩阵与向量的快速乘法（使用FFT）
        :param c: 循环矩阵第一行
        :param x: 向量
        :return: 乘积结果
        """
        # 使用FFT进行快速卷积
        n = len(c)
        c_fft = np.fft.fft(c)
        x_fft = np.fft.fft(x)
        result_fft = c_fft * x_fft
        return np.real(np.fft.ifft(result_fft))
    
    # ==================== Hankel矩阵 ====================
    
    @staticmethod
    def create_hankel_matrix(c: np.ndarray, r: Optional[np.ndarray] = None) -> np.ndarray:
        """
        创建Hankel矩阵
        :param c: 第一列
        :param r: 最后一行（可选）
        :return: Hankel矩阵
        """
        return hankel(c, r)
    
    # ==================== 范德蒙德矩阵 ====================
    
    @staticmethod
    def create_vandermonde_matrix(x: np.ndarray, n: Optional[int] = None) -> np.ndarray:
        """
        创建范德蒙德矩阵
        :param x: 基向量
        :param n: 列数（可选）
        :return: 范德蒙德矩阵
        """
        return np.vander(x, n)
    
    # ==================== 希尔伯特矩阵 ====================
    
    @staticmethod
    def create_hilbert_matrix(n: int) -> np.ndarray:
        """
        创建希尔伯特矩阵
        :param n: 矩阵维度
        :return: 希尔伯特矩阵
        """
        H = np.zeros((n, n))
        for i in range(n):
            for j in range(n):
                H[i, j] = 1.0 / (i + j + 1)
        return H
    
    @staticmethod
    def hilbert_matrix_inverse(n: int) -> np.ndarray:
        """
        希尔伯特矩阵的精确逆矩阵
        :param n: 矩阵维度
        :return: 希尔伯特矩阵的逆
        """
        from scipy.special import comb
        H_inv = np.zeros((n, n))
        
        for i in range(n):
            for j in range(n):
                sign = (-1) ** (i + j)
                H_inv[i, j] = sign * (i + j + 1) * comb(n + i, n - j - 1) * comb(n + j, n - i - 1) * comb(i + j, i) ** 2
        
        return H_inv
    
    # ==================== 随机矩阵 ====================
    
    @staticmethod
    def create_random_sparse_matrix(m: int, n: int, density: float = 0.1, 
                                   format: str = 'csr', random_state: Optional[int] = None) -> sparse.spmatrix:
        """
        创建随机稀疏矩阵
        :param m: 行数
        :param n: 列数
        :param density: 密度
        :param format: 稀疏矩阵格式
        :param random_state: 随机种子
        :return: 随机稀疏矩阵
        """
        return sparse.random(m, n, density=density, format=format, random_state=random_state)
    
    @staticmethod
    def create_random_symmetric_matrix(n: int, density: float = 0.1, 
                                      format: str = 'csr', random_state: Optional[int] = None) -> sparse.spmatrix:
        """
        创建随机对称稀疏矩阵
        :param n: 矩阵维度
        :param density: 密度
        :param format: 稀疏矩阵格式
        :param random_state: 随机种子
        :return: 随机对称稀疏矩阵
        """
        A = sparse.random(n, n, density=density, format='coo', random_state=random_state)
        A = A + A.T
        return A.asformat(format)
    
    # ==================== 矩阵结构检测 ====================
    
    @staticmethod
    def is_symmetric(matrix: Union[np.ndarray, sparse.spmatrix], tol: float = 1e-10) -> bool:
        """
        检测矩阵是否对称
        :param matrix: 输入矩阵
        :param tol: 容差
        :return: 是否对称
        """
        if sparse.issparse(matrix):
            return np.allclose((matrix - matrix.T).data, 0, atol=tol)
        else:
            return np.allclose(matrix, matrix.T, atol=tol)
    
    @staticmethod
    def is_hermitian(matrix: Union[np.ndarray, sparse.spmatrix], tol: float = 1e-10) -> bool:
        """
        检测矩阵是否厄米
        :param matrix: 输入矩阵
        :param tol: 容差
        :return: 是否厄米
        """
        if sparse.issparse(matrix):
            return np.allclose((matrix - matrix.H).data, 0, atol=tol)
        else:
            return np.allclose(matrix, np.conj(matrix.T), atol=tol)
    
    @staticmethod
    def is_positive_definite(matrix: np.ndarray) -> bool:
        """
        检测矩阵是否正定
        :param matrix: 输入矩阵
        :return: 是否正定
        """
        try:
            np.linalg.cholesky(matrix)
            return True
        except np.linalg.LinAlgError:
            return False
    
    @staticmethod
    def is_orthogonal(matrix: np.ndarray, tol: float = 1e-10) -> bool:
        """
        检测矩阵是否正交
        :param matrix: 输入矩阵
        :param tol: 容差
        :return: 是否正交
        """
        if matrix.shape[0] != matrix.shape[1]:
            return False
        
        product = matrix @ matrix.T
        identity = np.eye(matrix.shape[0])
        return np.allclose(product, identity, atol=tol)
    
    @staticmethod
    def matrix_structure_analysis(matrix: Union[np.ndarray, sparse.spmatrix]) -> dict:
        """
        矩阵结构分析
        :param matrix: 输入矩阵
        :return: 结构分析结果
        """
        result = {
            'shape': matrix.shape,
            'is_square': matrix.shape[0] == matrix.shape[1],
            'is_sparse': sparse.issparse(matrix)
        }
        
        if sparse.issparse(matrix):
            result.update(SpecialMatrices.sparse_matrix_stats(matrix))
        else:
            result.update({
                'dtype': matrix.dtype,
                'is_symmetric': SpecialMatrices.is_symmetric(matrix),
                'is_hermitian': SpecialMatrices.is_hermitian(matrix),
                'is_orthogonal': SpecialMatrices.is_orthogonal(matrix) if result['is_square'] else False,
                'is_positive_definite': SpecialMatrices.is_positive_definite(matrix) if result['is_square'] and SpecialMatrices.is_symmetric(matrix) else False
            })
        
        return result